"""LLM Plays Pokemon PyBoy模拟器实现模块"""

import os
import sys
import time
import logging
from typing import Dict, Any, List, Optional
from datetime import datetime

# 导入PyBoy
try:
    from pyboy import PyBoy
    # 定义按钮常量
    BUTTON_A = 0x01
    BUTTON_B = 0x02
    BUTTON_SELECT = 0x04
    BUTTON_START = 0x08
    BUTTON_RIGHT = 0x10
    BUTTON_LEFT = 0x20
    BUTTON_UP = 0x40
    BUTTON_DOWN = 0x80
except ImportError:
    logging.error(f"导入PyBoy库失败: {sys.exc_info()[1]}")
    sys.exit(1)  # 直接退出程序，返回错误码1

logger = logging.getLogger(__name__)


class PyBoyEmulator:
    """PyBoy模拟器实现类 - 简化版"""
    
    def __init__(self, config: Dict[str, Any]):
        """初始化PyBoy模拟器
        
        Args:
            config: 模拟器配置
        """
        self.config = config
        self.rom_path = config["rom_path"]
        self.save_state_dir = config["save_state_dir"]
        self.screenshot_dir = config["screenshot_dir"]
        self.pyboy_instance = None
        self.running = False
        
        # 按钮映射
        self.button_mapping = {
            "A": BUTTON_A,
            "B": BUTTON_B,
            "SELECT": BUTTON_SELECT,
            "START": BUTTON_START,
            "UP": BUTTON_UP,
            "DOWN": BUTTON_DOWN,
            "LEFT": BUTTON_LEFT,
            "RIGHT": BUTTON_RIGHT
        }
        
        # 确保目录存在
        os.makedirs(self.save_state_dir, exist_ok=True)
        os.makedirs(self.screenshot_dir, exist_ok=True)
        
        logger.info("PyBoy模拟器初始化完成")
    
    def start(self) -> bool:
        """启动模拟器并加载ROM
        
        Returns:
            是否成功启动
        """
        if self.running:
            logger.info("PyBoy模拟器已经在运行")
            return True
        
        logger.info("启动PyBoy模拟器并加载ROM...")
        
        try:
            # 检查ROM路径
            if not self.rom_path or not os.path.exists(self.rom_path):
                logger.error(f"ROM文件不存在: {self.rom_path}")
                return False
            
            # 启动PyBoy模拟器并直接加载ROM
            self.pyboy_instance = PyBoy(self.rom_path)
            self.running = True
            logger.info("PyBoy模拟器启动成功并已加载ROM")
            return True
        except Exception as e:
            logger.error(f"启动PyBoy模拟器失败: {e}")
            return False
    
    def stop(self) -> bool:
        """停止模拟器
        
        Returns:
            是否成功停止
        """
        if not self.running:
            logger.info("PyBoy模拟器未在运行")
            return True
        
        logger.info("停止PyBoy模拟器...")
        
        try:
            if self.pyboy_instance:
                self.pyboy_instance.stop(save=False)
                self.pyboy_instance = None
            
            self.running = False
            logger.info("PyBoy模拟器已停止")
            return True
        except Exception as e:
            logger.error(f"停止PyBoy模拟器失败: {e}")
            return False
    
    def save_state(self, slot: int = 0) -> bool:
        """保存游戏状态
        
        Args:
            slot: 存档槽位
            
        Returns:
            是否成功保存
        """
        if not self.running or not self.pyboy_instance:
            logger.error("PyBoy模拟器未启动，无法保存状态")
            return False
        
        save_path = os.path.join(self.save_state_dir, f"save_{slot}.state")
        
        try:
            self.pyboy_instance.save_state(save_path)
            logger.info(f"游戏状态已保存到: {save_path}")
            return True
        except Exception as e:
            logger.error(f"保存游戏状态失败: {e}")
            return False
    
    def load_state(self, slot: int = 0) -> bool:
        """加载游戏状态
        
        Args:
            slot: 存档槽位
            
        Returns:
            是否成功加载
        """
        if not self.running or not self.pyboy_instance:
            logger.error("PyBoy模拟器未启动，无法加载状态")
            return False
        
        save_path = os.path.join(self.save_state_dir, f"save_{slot}.state")
        
        if not os.path.exists(save_path):
            logger.error(f"存档文件不存在: {save_path}")
            return False
        
        try:
            self.pyboy_instance.load_state(save_path)
            logger.info(f"游戏状态已从 {save_path} 加载")
            return True
        except Exception as e:
            logger.error(f"加载游戏状态失败: {e}")
            return False
    
    def press_button(self, button: str, duration_ms: int = 100) -> bool:
        """按下按钮
        
        Args:
            button: 按钮名称
            duration_ms: 按下持续时间（毫秒）
            
        Returns:
            是否成功执行
        """
        if not self.running or not self.pyboy_instance:
            logger.error("PyBoy模拟器未启动，无法按下按钮")
            return False
        
        button = button.upper()
        
        if button not in self.button_mapping:
            logger.error(f"未知按钮: {button}")
            return False
        
        button_value = self.button_mapping[button]
        
        try:
            # 按下按钮
            self.pyboy_instance.send_input(button_value)
            # 等待指定时间
            time.sleep(duration_ms / 1000.0)
            # 释放按钮
            self.pyboy_instance.send_input(0)
            return True
        except Exception as e:
            logger.error(f"按下按钮失败: {e}")
            return False
    
    def press_buttons(self, buttons: List[str], pause_ms: Optional[int] = None) -> bool:
        """按下一系列按钮
        
        Args:
            buttons: 按钮序列
            pause_ms: 每个按钮之间的暂停时间（毫秒）
            
        Returns:
            是否成功执行
        """
        if not buttons:
            return True
        
        pause_ms = pause_ms or 100
        
        try:
            for button in buttons:
                if not self.press_button(button, pause_ms):
                    return False
                time.sleep(pause_ms / 1000.0)
            return True
        except Exception as e:
            logger.error(f"按下按钮序列失败: {e}")
            return False
    
    def take_screenshot(self) -> Dict[str, Any]:
        """获取游戏截图
        
        Returns:
            包含截图路径的字典
        """
        if not self.running or not self.pyboy_instance:
            return {
                "success": False,
                "message": "模拟器未启动",
                "screenshot": None
            }
        
        try:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            screenshot_path = os.path.join(self.screenshot_dir, f"screenshot_{timestamp}.png")
            
            # 获取屏幕图像并保存
            screen = self.pyboy_instance.screen_image()
            screen.save(screenshot_path)
            
            logger.info(f"截图已保存到: {screenshot_path}")
            
            return {
                "success": True,
                "message": "截图获取成功",
                "screenshot": screenshot_path
            }
        except Exception as e:
            logger.error(f"获取截图失败: {e}")
            return {
                "success": False,
                "message": f"获取截图失败: {e}",
                "screenshot": None
            }
    
    def tick(self, count: int = 1) -> bool:
        """执行游戏帧更新
        
        Args:
            count: 更新的帧数
            
        Returns:
            是否成功执行
        """
        if not self.running or not self.pyboy_instance:
            logger.error("PyBoy模拟器未启动，无法执行帧更新")
            return False
        
        try:
            for _ in range(count):
                self.pyboy_instance.tick()
            return True
        except Exception as e:
            logger.error(f"执行帧更新失败: {e}")
            return False
    
    def execute(self, buttons: List[str], pause_ms: Optional[int] = None, tick_count: int = 5) -> Dict[str, Any]:
        """执行按钮序列并更新游戏画面
        
        Args:
            buttons: 按钮序列
            pause_ms: 每个按钮之间的暂停时间（毫秒）
            tick_count: 执行的帧数
            
        Returns:
            执行结果，包含截图
        """
        if not self.running or not self.pyboy_instance:
            return {
                "success": False,
                "message": "模拟器未启动",
                "screenshot": None
            }
        
        # 按下按钮序列
        if not self.press_buttons(buttons, pause_ms):
            return {
                "success": False,
                "message": "按钮序列执行失败",
                "screenshot": None
            }
        
        # 更新游戏画面
        self.tick(tick_count)
        
        # 获取截图
        return self.take_screenshot()